#define K_MAX_SHAPE_DIM 0
#include "kernel_operator.h"
#include <type_traits>
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;
class KernelReduceLogSumExp_1 {
public:
    __aicore__ inline KernelReduceLogSumExp_1() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR axes, GM_ADDR y, 
                                int32_t repeat, int32_t cycles, TPipe* pipeIn) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x);
        yGm.SetGlobalBuffer((__gm__ DTYPE_X*)y);
        this->repeat = repeat;
        this->cycles = cycles;
        pipeIn->InitBuffer(inQueueX, BUFFER_NUM, (this->cycles+31)*sizeof(DTYPE_X));
        if constexpr (std::is_same_v<DTYPE_X, half>)
        {
            pipeIn->InitBuffer(QueueTmp1, (this->cycles+31) * sizeof(float));
        }
    }
    __aicore__ inline void Process() {
        for (int32_t i = 0; i < this->repeat; i++)
        {
            CopyIn(i*this->cycles, this->cycles);
            Compute(i, this->cycles);
        }
    }
private:
    __aicore__ inline void CopyIn(int32_t offset, int32_t lenght) {
        LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
        DataCopyExtParams copyParams{1, static_cast<uint32_t>(lenght * sizeof(DTYPE_X)), 0, 0, 0}; 
        DataCopyPadExtParams<DTYPE_X> padParams{false, 0, 0, 0};
        DataCopyPad(xLocal, xGm[offset], copyParams, padParams); 
        inQueueX.EnQue(xLocal);
    }
    __aicore__ inline void Compute(int32_t offset, int32_t lenght) {
        LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
        if constexpr (std::is_same_v<DTYPE_X, half>)
        {
            auto tmp1 = QueueTmp1.Get<float>();
            Cast(tmp1, xLocal, RoundMode::CAST_NONE, lenght);
            Exp(tmp1, tmp1, lenght);
            ReduceSum(tmp1, tmp1, tmp1, lenght);
            Ln(tmp1, tmp1, 1);
            Cast(xLocal, tmp1, RoundMode::CAST_NONE, 1);
            yGm.SetValue(offset, xLocal.GetValue(0));
        }
        else
        {
            Exp(xLocal, xLocal, lenght);
            ReduceSum(xLocal, xLocal, xLocal, lenght);
            Ln(xLocal, xLocal, 1);
            yGm.SetValue(offset, xLocal.GetValue(0));
        }
        inQueueX.FreeTensor(xLocal);
    }
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<QuePosition::VECCALC> QueueTmp1, QueueTmp2,boox;
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<DTYPE_X> yGm;
    int32_t repeat;
    int32_t cycles;
};
class KernelReduceLogSumExp {
public:
    __aicore__ inline KernelReduceLogSumExp() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR axes, GM_ADDR y, 
                                int32_t repeat, int32_t cycles, int32_t interval, 
                                int32_t loopCount, int32_t tileDataLenght, int32_t tailDataLenght, TPipe* pipeIn) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x);
        axesGm.SetGlobalBuffer((__gm__ int32_t*)axes);        
        yGm.SetGlobalBuffer((__gm__ DTYPE_X*)y);
        this->repeat = repeat;
        this->cycles = cycles;
        this->interval = interval;
        this->loopCount = loopCount;
        this->tileDataLenght = tileDataLenght;    
        this->tailDataLenght = tailDataLenght;
        pipeIn->InitBuffer(inQueueX, BUFFER_NUM, this->tileDataLenght*sizeof(DTYPE_X));
        pipeIn->InitBuffer(outQueueY, BUFFER_NUM, this->tileDataLenght*sizeof(DTYPE_X));
        if constexpr (std::is_same_v<DTYPE_X, half>)
        {
            pipeIn->InitBuffer(QueueTmp1, this->tileDataLenght * sizeof(float));
            pipeIn->InitBuffer(QueueTmp2, this->tileDataLenght * sizeof(float));
        }
    }
    __aicore__ inline void Process() {
        for (int32_t i = 0; i < this->repeat; i++) 
        {
            if constexpr (std::is_same_v<DTYPE_X, half>)
            {
                auto tmp2 = QueueTmp2.Get<float>();
                this->processDataLenght = this->tileDataLenght;
                for(int32_t j=0; j<this->loopCount; j++)
                {
                    Duplicate(tmp2, (float)0, this->processDataLenght);
                    for(int32_t z = 0; z < this->cycles; z++)
                    {
                        CopyIn(i, j, z);
                        Compute(tmp2);
                    }
                    LocalTensor<DTYPE_X> yLocal = outQueueY.AllocTensor<DTYPE_X>();
                    Ln(tmp2, tmp2, this->processDataLenght);
                    Cast(yLocal, tmp2, RoundMode::CAST_NONE, this->processDataLenght);
                    outQueueY.EnQue(yLocal);
                    CopyOut(i, j);
                }
                this->processDataLenght = this->tailDataLenght;
                Duplicate(tmp2, (float)0, this->processDataLenght);
                for(int32_t z = 0; z < this->cycles; z++)
                {
                    CopyIn(i, this->loopCount,z);
                    Compute(tmp2);
                }
                LocalTensor<DTYPE_X> yLocal = outQueueY.AllocTensor<DTYPE_X>();
                Ln(tmp2, tmp2, this->processDataLenght);
                Cast(yLocal, tmp2, RoundMode::CAST_NONE, this->processDataLenght);
                outQueueY.EnQue(yLocal);
                CopyOut(i, this->loopCount);
            }
            else
            {
                this->processDataLenght = this->tileDataLenght;
                for(int32_t j=0; j<this->loopCount; j++)
                {
                    LocalTensor<float> yLocal = outQueueY.AllocTensor<float>();
                    Duplicate(yLocal, (float)0, this->processDataLenght);
                    for(int32_t z = 0; z < this->cycles; z++)
                    {
                        CopyIn(i, j, z);
                        Compute(yLocal);
                    }
                    Ln(yLocal, yLocal, this->processDataLenght);
                    outQueueY.EnQue(yLocal);
                    CopyOut(i, j);
                }
                this->processDataLenght = this->tailDataLenght;
                LocalTensor<float> yLocal = outQueueY.AllocTensor<float>();
                Duplicate(yLocal, (float)0, this->processDataLenght);
                for(int32_t z = 0; z < this->cycles; z++)
                {
                    CopyIn(i, this->loopCount, z);
                    Compute(yLocal);
                }
                Ln(yLocal, yLocal, this->processDataLenght);
                outQueueY.EnQue(yLocal);
                CopyOut(i, this->loopCount);
            }
        }
    }
private:
    __aicore__ inline void CopyIn(int32_t offset, int32_t step, int32_t cycle) {
        LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
        DataCopyExtParams copyParams{1, static_cast<uint32_t>(this->processDataLenght * sizeof(DTYPE_X)), 0, 0, 0}; 
        DataCopyPadExtParams<DTYPE_X> padParams{false, 0, 0, 0};
        DataCopyPad(xLocal, xGm[offset*this->cycles*this->interval + step*this->tileDataLenght +cycle*this->interval], copyParams, padParams); 
        inQueueX.EnQue(xLocal);
    }
    __aicore__ inline void Compute(LocalTensor<float> yLocal) {
        LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
        if constexpr (std::is_same_v<DTYPE_X, half>)
        {
            auto tmp1 = QueueTmp1.Get<float>();
            Cast(tmp1, xLocal, RoundMode::CAST_NONE, this->processDataLenght);
            Exp(tmp1, tmp1, this->processDataLenght);
            Add(yLocal, yLocal, tmp1, this->processDataLenght);
        }
        else
        {
            Exp(xLocal, xLocal, this->processDataLenght);
            Add(yLocal, yLocal, xLocal.ReinterpretCast<float>(), this->processDataLenght);
        }
        inQueueX.FreeTensor(xLocal);
    }
    __aicore__ inline void CopyOut(int32_t progress, int32_t step) {
        LocalTensor<DTYPE_X> yLocal = outQueueY.DeQue<DTYPE_X>();
        DataCopyExtParams copyParams{1, static_cast<uint32_t>(this->processDataLenght * sizeof(DTYPE_X)), 0, 0, 0}; 
        DataCopyPad(yGm[progress*this->interval + step*this->tileDataLenght], yLocal, copyParams);
        outQueueY.FreeTensor(yLocal);
    }
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<QuePosition::VECCALC> QueueTmp1, QueueTmp2,boox;
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<int32_t> axesGm;
    GlobalTensor<DTYPE_X> yGm;
    int32_t cycles;
    int32_t interval;
    int32_t repeat;
    uint32_t loopCount;
    uint32_t tileDataLenght;
    uint32_t tailDataLenght;
    uint32_t processDataLenght;
};
class KernelReduceLogSumExp2 {
public:
    __aicore__ inline KernelReduceLogSumExp2() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR axes, GM_ADDR y, 
                                int32_t repeat, int32_t cycles, int32_t interval, 
                                int32_t loopCount, int32_t tileDataLenght, int32_t tailDataLenght, TPipe* pipeIn) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x);
        axesGm.SetGlobalBuffer((__gm__ int32_t*)axes);        
        yGm.SetGlobalBuffer((__gm__ DTYPE_X*)y);
        this->repeat = repeat;
        this->cycles = cycles;
        this->interval = interval;
        this->loopCount = loopCount;
        this->tileDataLenght = tileDataLenght;    
        this->tailDataLenght = tailDataLenght;
        pipeIn->InitBuffer(inQueueX, BUFFER_NUM, this->tileDataLenght*this->interval*sizeof(DTYPE_X));
        pipeIn->InitBuffer(outQueueY, BUFFER_NUM, this->tileDataLenght*this->interval*sizeof(DTYPE_X));
        if constexpr (std::is_same_v<DTYPE_X, half>)
        {
            pipeIn->InitBuffer(QueueTmp1, this->tileDataLenght*this->interval * sizeof(float));
            pipeIn->InitBuffer(QueueTmp2, this->tileDataLenght*this->interval * sizeof(float));
        }
    }
    __aicore__ inline void Process() {
        if(this->loopCount != 0)
        {
        for (int32_t i = 0; i < this->repeat; i++) 
        {
            if constexpr (std::is_same_v<DTYPE_X, half>)
            {
                auto tmp2 = QueueTmp2.Get<float>();
                this->processDataLenght = this->tileDataLenght;
                Duplicate(tmp2, (float)0, this->processDataLenght*this->interval);
                for(int32_t j=0; j<this->loopCount; j++)
                {
                    CopyIn(i, j);
                    Compute(tmp2, this->processDataLenght*this->interval);
                }
                this->processDataLenght = this->tailDataLenght;
                CopyIn(i, this->loopCount);
                Compute(tmp2, this->processDataLenght*this->interval);
                LocalTensor<DTYPE_X> yLocal = outQueueY.AllocTensor<DTYPE_X>();
                for(int i=1; i<this->tileDataLenght;i++)
                {
                    Add(tmp2, tmp2, tmp2[i*this->interval], this->interval);
                }
                Ln(tmp2, tmp2, this->interval);
                Cast(yLocal, tmp2, RoundMode::CAST_NONE, this->interval);
                outQueueY.EnQue(yLocal);
                CopyOut(i);
            }
            else
            {
                this->processDataLenght = this->tileDataLenght;
                LocalTensor<float> yLocal = outQueueY.AllocTensor<float>();
                Duplicate(yLocal, (float)0, this->processDataLenght*this->interval);
                for(int32_t j=0; j<this->loopCount; j++)
                {
                    CopyIn(i, j);
                    Compute(yLocal, this->processDataLenght*this->interval);
                }
                this->processDataLenght = this->tailDataLenght;
                CopyIn(i, this->loopCount);
                Compute(yLocal, this->processDataLenght*this->interval);
                for(int i=1; i<this->tileDataLenght;i++)
                {
                    Add(yLocal, yLocal, yLocal[i*this->interval], this->interval);
                }
                Ln(yLocal, yLocal, this->interval);
                outQueueY.EnQue(yLocal);
                CopyOut(i);
            }
        }
        }
        else
        {
            for (int32_t i = 0; i < this->repeat; i++) 
            {
                if constexpr (std::is_same_v<DTYPE_X, half>)
                {
                    auto tmp2 = QueueTmp2.Get<float>();
                    this->processDataLenght = this->tailDataLenght;
                    Duplicate(tmp2, (float)0, this->processDataLenght*this->interval);
                    CopyIn(i, this->loopCount);
                    Compute(tmp2, this->processDataLenght*this->interval);
                    LocalTensor<DTYPE_X> yLocal = outQueueY.AllocTensor<DTYPE_X>();
                    for(int i=1; i<this->tailDataLenght;i++)
                    {
                        Add(tmp2, tmp2, tmp2[i*this->interval], this->interval);
                    }
                    Ln(tmp2, tmp2, this->interval);
                    Cast(yLocal, tmp2, RoundMode::CAST_NONE, this->interval);
                    outQueueY.EnQue(yLocal);
                    CopyOut(i);
                }
                else
                {
                    this->processDataLenght = this->tailDataLenght;
                    LocalTensor<float> yLocal = outQueueY.AllocTensor<float>();
                    Duplicate(yLocal, (float)0, this->processDataLenght*this->interval);
                    CopyIn(i, this->loopCount);
                    Compute(yLocal, this->processDataLenght*this->interval);
                    for(int i=1; i<this->tailDataLenght;i++)
                    {
                        Add(yLocal, yLocal, yLocal[i*this->interval], this->interval);
                    }
                    Ln(yLocal, yLocal, this->interval);
                    outQueueY.EnQue(yLocal);
                    CopyOut(i);
                }
            }
        }
    }
private:
    __aicore__ inline void CopyIn(int32_t offset, int32_t step) {
        LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
        DataCopy(xLocal, xGm[offset*this->cycles*this->interval + step*this->tileDataLenght*this->interval], this->processDataLenght*this->interval); 
        inQueueX.EnQue(xLocal);
    }
    __aicore__ inline void Compute(LocalTensor<float> yLocal, uint32_t lenght) {
        LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
        if constexpr (std::is_same_v<DTYPE_X, half>)
        {
            auto tmp1 = QueueTmp1.Get<float>();
            Cast(tmp1, xLocal, RoundMode::CAST_NONE, lenght);
            Exp(tmp1, tmp1, lenght);
            Add(yLocal, yLocal, tmp1, lenght);
        }
        else
        {
            Exp(xLocal, xLocal, lenght);
            Add(yLocal, yLocal, xLocal.ReinterpretCast<float>(), lenght);
        }
        inQueueX.FreeTensor(xLocal);
    }
    __aicore__ inline void CopyOut(int32_t progress) {
        LocalTensor<DTYPE_X> yLocal = outQueueY.DeQue<DTYPE_X>();
        DataCopy(yGm[progress*this->interval], yLocal, this->interval);
        outQueueY.FreeTensor(yLocal);
    }
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<QuePosition::VECCALC> QueueTmp1, QueueTmp2,boox;
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<int32_t> axesGm;
    GlobalTensor<DTYPE_X> yGm;
    int32_t cycles;
    int32_t interval;
    int32_t repeat;
    uint32_t loopCount;
    uint32_t tileDataLenght;
    uint32_t tailDataLenght;
    uint32_t processDataLenght;
};
extern "C" __global__ __aicore__ void reduce_log_sum_exp(GM_ADDR x, GM_ADDR axes, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    TPipe pipe;
    if(TILING_KEY_IS(1)) { 
        KernelReduceLogSumExp_1 op;
        op.Init(x, axes, y, 
            tiling_data.repeat, tiling_data.cycles, &pipe);
        op.Process();
    }
    else if(TILING_KEY_IS(2)) { 
        KernelReduceLogSumExp op;
        op.Init(x, axes, y, 
                tiling_data.repeat, tiling_data.cycles, tiling_data.interval, 
                tiling_data.loopCount, tiling_data.tileDataLenght, tiling_data.tailDataLenght, &pipe);
        op.Process();
    }
    else if(TILING_KEY_IS(3)) { 
        KernelReduceLogSumExp2 op;
        op.Init(x, axes, y, 
                tiling_data.repeat, tiling_data.cycles, tiling_data.interval, 
                tiling_data.loopCount, tiling_data.tileDataLenght, tiling_data.tailDataLenght, &pipe);
        op.Process();
    }
}